import yaml
from MLModels import *
import os
from multiprocessing import cpu_count 

yaml_file = "/home/cyw/projects/malware_detected/configue.yaml"
with open(yaml_file, 'r') as f:
    config = yaml.safe_load(f)
debugLog = config["logs"]["debugLog"]

class modelFactory(modelTrain):
    def __init__(self,modelName=None):
        if modelName==None:
            modelName = config["model"]["trainModel"]
        print("使用工厂方法加载{}模型".format(modelName))
        if modelName == "lightGbm":
            self.model = lightGbmTrain()
        elif modelName == "decisionTree":
            self.model = decisionTreeTrain()
        elif modelName == "xgBoost":
            self.model = xgBoostTrain()
        elif modelName == "knn":
            self.model = KNNTrain()
        elif modelName == "naiveBayse":
            self.model = NBTrain()
        elif modelName == "randomForest":
            self.model = rfTrain()
        elif modelName == "svm":
            self.model = svmTrain()
        elif modelName == "dnn":
            self.model = DNNTrain()
        else:
            assert False, "{}模型并未实现".format(modelName)

    def trainDataSet(self):
        """
            训练集数据加载
        """
        return self.model.trainDataSet()
        
    
    def train(self):
        """
            训练模型
        """
        self.model.train(self.trainDataSet())
    
if __name__ == "__main__":
    a = time.time()
    trainTool = modelFactory()
    trainTool.train()
    b = time.time()
    print("模型训练耗时{}".format(b-a))